import { BaseAgentBackend, AgentBackendFactory } from "../../../public/agentBackend";
import { getEnv, EnvEnum } from "../../../public/env";
import { eclConfig } from "../config";

export interface EmbeddingConfig {
    /** 模型制造商 */
    modelMaker: string;
    /** 嵌入模型类型 */
    embeddingModelType: string;
    /** 单次最大token数，默认 8191 */
    maxTokens?: number;
    /** 累计token数 */
    promptTokenCount?: number;
    /** 总token数 */
    propmtTotalCount?: number;
}

export class EmbeddingBackend {
    public modelMaker: string;
    public embeddingModelType: string;
    public maxTokens: number;

    private key: string;
    private modelBackend: typeof BaseAgentBackend;

    constructor(config: EmbeddingConfig) {
        this.modelMaker = config.modelMaker;
        this.embeddingModelType = config.embeddingModelType;
        this.maxTokens = config?.maxTokens ?? 8191;

        this.key = getEnv(EnvEnum.AGENT_BACKEND_API_KEY);
        this.modelBackend = AgentBackendFactory.create(this.modelMaker);
    }

    async getEmbedding(content: string): Promise<number[]> {
        const validatedContent = this.processContent(content);
        return await this.request(validatedContent);
    }

    private async request(content: string, retryCount: number = 0): Promise<number[]> {
        try {
            const response = await this.modelBackend.getTextEmbedding({
                apiKey: this.key,
                input: content,
                model: this.embeddingModelType
            });
            const embedding = response.data[0].embedding;
            if (!embedding) {
                throw new Error("embedding 为空");
            }
            return embedding;
        } catch (error) {
            if (retryCount < eclConfig.retryNum) {
                return await this.request(content, retryCount + 1);
            }
            throw new Error(
                `无法获取 embedding 结果，请检查网络连接或重试: ${(error as Error).message}`
            );
        }
    }

    private processContent(content: string): string {
        if (content.length > this.maxTokens) {
            return content.substring(0, this.maxTokens);
        }
        return content;
    }
}
